#include <windows.h>		// temp

#include "zlib.h"

#include "lib.h"
#include "stream.h"
#include "endian.h"

#include "zipfile.h"

#include "pack8.h"

// "ZFH" means Zip File Header -- a header of some sort
//
class ZFHCenDirEnd			// CENtral DIRectory END
{
	public:
		ZFHCenDirEnd() { pos = -1; comment = (uint8*)0; }
		~ZFHCenDirEnd() { delete [] comment; }

		bool8 load(ZipFile*);

	protected:
		int32 locate(ZipFile*);

	public:
		enum Constants
		{
			CENDIREND_SIZE	= 4+2+2+2+2+4+4+2,
			MAXCDIR_SIZE	= 0x10000L
#	define ZFHCenDirEnd_Endian "w4w2w2w2w2w4w4w2"
		};

		uint32	signature;				// 0x06054b50
		uint16	diskNum;				// this disk's number
		uint16	diskNumCenDirStart;		// disk # that has start of CenDir
		uint16	diskNumEntries;			// # zip entries on this disk
		uint16  totalEntries;			// # zip entries in entire zipfile
		uint32	cenDirSize;				// total size of CenDir
		uint32	cenDirOfs;				// ofs on disk diskNumCenDirStart
		uint16	commentLen;
		uint8*	comment;

		int32 pos;
};

class ZFHEntry
{
	public:
//		ZFHEntry() { filename = (uint8*)0; }
		ZFHEntry() { /*filename = "";*/ }
		~ZFHEntry() { /*delete [] filename;*/ }

		bool8 load(ZipFile*);

	public:
		enum Constants
		{
			ZIPENTRY_SIZE	= 4+2+2+2+2+2+2+4+4+4+2+2+2+2+2+4+4,
			STORED			= 0,
			DEFLATED		= 8
#	define ZFHEntry_Endian	"w4w2w2w2w2w2w2w4w4w4w2w2w2w2w2w4w4"
		};

		uint32 signature;				// 0x02014b50
		uint16 version;					// zip version of maker
		uint16 extractVer;				// zip version needed to extract
		uint16 flags;
		uint16 method;					// compression method
		uint16 lastModTime;				// MSDOS format
		uint16 lastModDate;				// MSDOS format
		uint32 filecrc32;
		uint32 compressedSize;
		uint32 uncompressedSize;
		uint16 filenameLen;
		uint16 extraFieldLen;
		uint16 fileCommentLen;
		uint16 diskNumStart;			// disk # that starts this entry's data
		uint16 internalAttr;
		uint32 externalAttr;
		uint32 headerOfs;
//		uint8* filename;
		String filename;
		// extra and comment are ignored
};

class ZFHLocalEntry
{
	public:
		ZFHLocalEntry() { }

		bool8 load(ZipFile* ZF);

	public:
		enum Constants
		{
			LOCALENTRY_SIZE = 4+2+2+2+2+2+4+4+4+2+2
#	define ZFHLocalEntry_Endian	"w4w2w2w2w2w2w4w4w4w2w2"
		};

		uint32 signature;               // 0x04034b50
		uint16 extractVer;              // zip version needed to extract
		uint16 flags;
		uint16 method;                  // compression method
		uint16 lastModTime;             // MSDOS format
		uint16 lastModDate;             // MSDOS format
		uint32 filecrc32;
		uint32 compressedSize;
		uint32 uncompressedSize;
		uint16 filenameLen;
		uint16 extraFieldLen;
		// extra and filename are ignored
};

#include "packend.h"

class inflateStream : public UnzipStream
{
	public:
		inflateStream(ZipFile* ZF, ZipEntry* ZE);
		~inflateStream();

		int32 read(uint8* buf, int32 bufSize);

	protected:
		enum Constants
		{
			FBUFSIZE = 4096,
			ZSTREAM_NOTINIT = 0x1L,
			ZSTREAM_EOFSEEN = 0x2L
		};
		z_stream zstream;
		int32 flags;
		uint8 fbuf[FBUFSIZE];
};

class unstoreStream : public UnzipStream
{
	public:
		unstoreStream(ZipFile* ZF, ZipEntry* ZE);

		int32 read(uint8* buf, int32 bufSize);
};

//
// -------------------------- ZipFile
//

ZipFile::ZipFile()
{
	file = (pCStream)0;
	flen = 0;
	entry = (ZipEntry*)0;
}

ZipFile::~ZipFile()
{
	close();
}

bool8
ZipFile::open(const String& filename)
{
	ZFHCenDirEnd cde;
	ZFHEntry zfhe;
	ZipEntry* ze = (ZipEntry*)0;
	ZipEntry* lastZE = (ZipEntry*)0;

	close();

	file = newCStream(filename, streamO_RDONLY);
	if (!file)
		goto abort;

	if (!determineFileLen())
		goto abort;

	if (!cde.load(this))
		goto abort;
	
	if (streamEOF == seek(cde.cenDirOfs, streamSEEK_SET))
		goto abort;

	while (tell() < cde.pos)
	{
		if (!zfhe.load(this))
			goto abort;

		ze = new ZipEntry(zfhe);
		if (!ze)
			goto abort;

		if (lastZE)
			lastZE->next = ze;
		else
			entry = ze;
		lastZE = ze;
	}

	return true;

abort:
	close();
	return false;
}

void
ZipFile::close()
{
	ZipEntry* temp;

	deleteCStream(file);
	file = (pCStream)0;
	flen = 0;

	while (entry)
	{
		temp = entry;
		entry = entry->next;
		delete temp;
	}
}

bool8
ZipFile::determineFileLen()
{
	if (streamEOF == seek(0, streamSEEK_END))
		return false;
	flen = tell();
	if (flen == streamEOF)
		return false;

	// don't seek back to beginning...first action after this method is to
	// search for CenDirEnd from end of the file.
	//
	return true;
}

int32
ZipFile::read(void* Buf, int32 Amt)
{
	return CStreamRead(file, (uint8*)Buf, Amt);
}

int32
ZipFile::seek(int32 Ofs, int32 Whence)
{
	return CStreamSeek(file, Ofs, Whence);
}

int32
ZipFile::tell()
{
	return CStreamTell(file);
}

//
// -------------------------- ZFHCenDirEnd
//

bool8
ZFHCenDirEnd::load(ZipFile* ZF)
{
	pos = locate(ZF);

	if (pos >= 0 &&
		streamEOF != ZF->seek(pos, streamSEEK_SET) &&
		CENDIREND_SIZE == ZF->read(&signature, CENDIREND_SIZE))
	{
		endianL2M(&signature, ZFHCenDirEnd_Endian);

		if (commentLen > 0)
		{
			comment = new uint8[commentLen + 1];
			if (comment)
			{
				if (commentLen != ZF->read(comment, commentLen))
					commentLen = 0;

				comment[commentLen] = '\0';
			}
		}
		else
			comment = (uint8*)0;

		if (diskNum != diskNumCenDirStart)	// disk-spanning not implemented
			return false;

		return true;
	}
	return false;
}

int32
ZFHCenDirEnd::locate(ZipFile* ZF)
{
	int32 pos, amt, searched, i;
	const int32 BUFLEN = 4096;
	uint8 buf[BUFLEN];

	// search for end of central directory
	//
	amt = MIN(ZF->getFileLen(), BUFLEN);
	if (amt < 4)
		return -1;		// file too small (way!)

	pos = ZF->getFileLen() - amt;
	searched = 0;
	while (pos >= 0 && searched < MAXCDIR_SIZE)
	{
		if (streamEOF == ZF->seek(pos, streamSEEK_SET))
			return -1;

		if (amt != ZF->read(buf, amt))
			return -1;

		// search backwards for end of central directory
		//
		for (i = amt - 4; i >= 0; i--)
			if (buf[i+0] == 0x50 && 
				buf[i+1] == 0x4b &&
				buf[i+2] == 0x05 &&
				buf[i+3] == 0x06 &&
				ZF->getFileLen() - (pos + i) >= CENDIREND_SIZE)
				return pos + i;

		searched = ZF->getFileLen() - pos;
		amt = MIN(ZF->getFileLen() - searched, BUFLEN);
		if (amt == 0)						// searched entire file ( <64K )
			break;
		pos -= amt;
		if (amt > BUFLEN - 3)				// overlap 3 bytes (4 byte pattern)
			pos += 3;
		else								// final search, overlap with prev
			amt = MIN(amt + 3, ZF->getFileLen());
	}

	return -1;
}

//
// -------------------------- ZFHEntry
//

bool8
ZFHEntry::load(ZipFile* ZF)
{
	int32 pos;
	uint8* name;

	// assumed: file pointer positioned at start of next ZFHEntry to read
	//
	pos = ZF->tell();
	if (pos == streamEOF)
		return false;

	if (ZIPENTRY_SIZE == ZF->read(&signature, ZIPENTRY_SIZE))
	{
		endianL2M(&signature, ZFHEntry_Endian);

		if (filenameLen > 0)
		{
			name = new uint8[filenameLen + 1];
			if (name)
			{
				if (filenameLen != ZF->read(name, filenameLen))
					filenameLen = 0;

				name[filenameLen] = '\0';
				filename = (char*)name;
				delete name;
			}
		}
		else
			name = (uint8*)0;

		// skip the extra comment fields
		//
		if (streamEOF == ZF->seek(
			pos + ZIPENTRY_SIZE + filenameLen + extraFieldLen + fileCommentLen,
			streamSEEK_SET))
			return false;

		if (signature != 0x02014b50 ||
			!(method == STORED || method == DEFLATED))
			return false;

		return true;
	}
	return false;
}

//
// -------------------------- ZFHLocalEntry
//

bool8
ZFHLocalEntry::load(ZipFile* ZF)
{
	// assumed: file pointer positioned at start of next ZFHLocalEntry to read
	//
	if (LOCALENTRY_SIZE == ZF->read(&signature, LOCALENTRY_SIZE))
	{
		endianL2M(&signature, ZFHLocalEntry_Endian);

		// skip the filename & extra fields
		//
		if (streamEOF == ZF->seek(filenameLen + extraFieldLen, streamSEEK_CUR))
			return false;

		if (signature != 0x04034b50)
			return false;

		return true;
	}
	return false;
}

//
// -------------------------- ZipEntry
//

ZipEntry::ZipEntry(ZFHEntry& zfhe)
{
	flags = zfhe.flags;
	method = zfhe.method;
	filecrc32 = zfhe.filecrc32;
	compressedSize = zfhe.compressedSize;
	uncompressedSize = zfhe.uncompressedSize;
	headerOfs = zfhe.headerOfs;
	filename = zfhe.filename;
	next = (ZipEntry*)0;
}

bool8
ZipEntry::isDirectory()
{
	return (bool8)(filename[ filename.getLength() ] == '/');
}

UnzipStream*
ZipEntry::getUnzipStream(ZipFile* ZF)
{
	ZFHLocalEntry zel;

	if (isDirectory() || isEncrypted())
		return (UnzipStream*)0;

	if (streamEOF == ZF->seek(headerOfs, streamSEEK_SET))
		return (UnzipStream*)0;

	if (!zel.load(ZF))
		return (UnzipStream*)0;

	if (method == DEFLATED)
		return new inflateStream(ZF, this);
	else if (method == STORED)
		return new unstoreStream(ZF, this);

	return (UnzipStream*)0;
}

//
// -------------------------- UnzipStream
//

UnzipStream::UnzipStream(ZipFile* ZF, ZipEntry* ZE)
{
	zf = ZF;
	ze = ZE;
	remain = ze->compressedSize;
	crc = crc32(0L, Z_NULL, 0);
}

inflateStream::inflateStream(ZipFile* ZF, ZipEntry* ZE)
	: UnzipStream(ZF, ZE)
{
	uint8* p = (uint8*)&zstream;
	for (int32 i = 0; i < sizeof(z_stream); i++)
		*p++ = 0;

	// zlib normally requires its own headers and trailers around the
	// compressed data; however, there is an undocumented way to turn off
	// the zlib headers/trailers by passing a -bits to inflateInit2_.
	// After doing so, zlib expects raw compressed data + ONE DUMMY BYTE
	// (It only requires the dummy byte if we expect inflate() to eventually
	// return Z_STREAM_END).  So, we read one extra byte.
	//
	++remain;

	flags = 0;

	if (Z_OK != inflateInit2_(&zstream, -MAX_WBITS,
		ZLIB_VERSION, sizeof(z_stream)))
		flags |= ZSTREAM_NOTINIT;
}

inflateStream::~inflateStream()
{
	if ((flags & ZSTREAM_NOTINIT) == 0)
		inflateEnd(&zstream);
}

int32
inflateStream::read(uint8* buf, int32 bufSize)
{
	int32 amt;
	int err;

	if (flags == 0)
	{
		if (zstream.avail_in == 0)
		{
			amt = MIN(remain, FBUFSIZE);
			amt = zf->read(fbuf, amt);
			if (amt == streamEOF)
				return -2;
			remain -= amt;

			zstream.next_in  = (Bytef*)fbuf;
			zstream.avail_in = (uInt)amt;
		}

		zstream.next_out = (uint8*)buf;
		zstream.avail_out = (uInt)bufSize;

		err = inflate(&zstream, Z_NO_FLUSH);
		if (err != Z_STREAM_END && err != Z_OK)
			return -2;

		amt = bufSize - zstream.avail_out;
		if (amt)
			crc = crc32(crc, buf, amt);

		if (err == Z_STREAM_END)
			flags |= ZSTREAM_EOFSEEN;

		return amt;
	}

	if (flags & ZSTREAM_EOFSEEN)
	{
		flags &= ~ZSTREAM_EOFSEEN;
		flags |= ZSTREAM_NOTINIT;

		if (Z_OK == inflateEnd(&zstream))
		{
			if (ze->flags & 8)	// read trailing data descriptor if it exists
			{
				struct descrip { uint32 l[4]; } dd;
 
				// read one too far (gave zlib a dummy byte)
				//
				if (streamEOF != zf->seek(-1, streamSEEK_CUR) &&
					16 == zf->read(&dd, 16) &&
					dd.l[0] == 0x08074b50)
				{
					ze->filecrc32 = dd.l[1];
					ze->compressedSize = dd.l[2];
					ze->uncompressedSize = dd.l[3];
				}
				else
					return -2;
			}
			
			if (crc == ze->filecrc32)
				return -1;
		}
	}

	return -2;
}

unstoreStream::unstoreStream(ZipFile* ZF, ZipEntry* ZE)
	: UnzipStream(ZF, ZE)
{
	remain = ze->uncompressedSize;
}

int32
unstoreStream::read(uint8* buf, int32 bufSize)
{
	int32 amt = MIN(remain, bufSize);

	amt = zf->read(buf, amt);
	if (amt != streamEOF)
	{
		crc = crc32(crc, buf, amt);
		remain -= amt;
	}
	else
	{
		if (crc != ze->filecrc32)
			amt = -2;
	}

	return amt;
}
